import gym
import numpy as np
from gym import spaces


class PVZEnv(gym.Env):
    """植物大战僵尸游戏环境"""

    def __init__(self):
        super(PVZEnv, self).__init__()

        # 定义游戏参数
        self.num_rows = 5  # 草坪行数
        self.num_cols = 9  # 草坪列数
        self.max_plants = 5  # 每种植物的最大数量
        self.max_zombies = 10  # 最大僵尸数量
        self.sun_init = 50  # 初始阳光数量

        # 定义植物类型和属性 (类型, 花费, 生命值, 攻击力)
        self.plants = {
            0: ('豌豆射手', 100, 100, 20),
            1: ('向日葵', 50, 80, 0),
            2: ('坚果墙', 50, 400, 0),
            3: ('寒冰射手', 175, 100, 20),
            4: ('双发射手', 200, 100, 40),
            5: ('大嘴花', 150, 150, 30),
            6: ('火爆辣椒', 125, 100, 100),
            7: ('土豆雷', 25, 50, 150),
            8: ('食人花', 150, 150, 40),
            9: ('三线射手', 325, 100, 60)
        }

        # 定义僵尸类型和属性 (类型, 生命值, 攻击力, 速度)
        self.zombies = {
            0: ('普通僵尸', 200, 10, 1),
            1: ('路障僵尸', 400, 10, 1)
        }

        # 定义动作空间: (行, 列, 植物类型)
        self.action_space = spaces.MultiDiscrete([self.num_rows, self.num_cols, len(self.plants)])

        # 定义观察空间
        # 包括：草坪状态(行×列×(植物类型+僵尸类型))、阳光数量、各类植物剩余数量
        obs_shape = (
            self.num_rows,
            self.num_cols,
            len(self.plants) + len(self.zombies)
        )
        self.observation_space = spaces.Dict({
            'lawn': spaces.Box(low=0, high=1, shape=obs_shape, dtype=np.float32),
            'sun': spaces.Box(low=0, high=float('inf'), shape=(1,), dtype=np.float32),
            'plants_remaining': spaces.Box(low=0, high=self.max_plants, shape=(len(self.plants),), dtype=np.int32)
        })

        # 初始化游戏状态
        self.reset()

    def reset(self):
        """重置环境状态"""
        self.lawn = np.zeros((self.num_rows, self.num_cols, len(self.plants) + len(self.zombies)))
        self.sun = self.sun_init
        self.plants_remaining = np.full(len(self.plants), self.max_plants)
        self.zombies_remaining = self.max_zombies
        self.game_over = False
        self.steps = 0
        return self._get_obs()

    def step(self, action):
        """执行一步动作"""
        reward = 0
        self.steps += 1

        # 解析动作
        row, col, plant_type = action

        # 检查动作是否有效
        if not self._is_valid_action(row, col, plant_type):
            return self._get_obs(), -10, True, {}

        # 放置植物
        self._place_plant(row, col, plant_type)

        # 生成僵尸
        if self.steps % 5 == 0 and self.zombies_remaining > 0:
            self._spawn_zombie()

        # 更新游戏状态
        reward += self._update_state()

        # 检查游戏是否结束
        done = self.game_over or self.zombies_remaining <= 0

        return self._get_obs(), reward, done, {}

    def _get_obs(self):
        """获取当前观察状态"""
        return {
            'lawn': self.lawn.astype(np.float32),
            'sun': np.array([self.sun], dtype=np.float32),
            'plants_remaining': self.plants_remaining
        }

    def _is_valid_action(self, row, col, plant_type):
        """检查动作是否有效"""
        # 检查位置是否已被占用
        if np.sum(self.lawn[row, col]) > 0:
            return False

        # 检查是否有足够的阳光
        if self.sun < self.plants[plant_type][1]:
            return False

        # 检查是否还有剩余植物
        if self.plants_remaining[plant_type] <= 0:
            return False

        return True

    def _place_plant(self, row, col, plant_type):
        """放置植物"""
        self.lawn[row, col, plant_type] = 1
        self.sun -= self.plants[plant_type][1]
        self.plants_remaining[plant_type] -= 1

    def _spawn_zombie(self):
        """生成僵尸"""
        row = np.random.randint(0, self.num_rows)
        zombie_type = np.random.randint(0, len(self.zombies))
        # 移除对该位置是否被占用的检查，允许在同一位置生成多个僵尸
        self.lawn[row, -1, len(self.plants) + zombie_type] = 1
        self.zombies_remaining -= 1

    def _update_state(self):
        """更新游戏状态"""
        reward = 0

        # 向日葵产生阳光
        sunflowers = np.sum(self.lawn[:, :, 1])
        self.sun += sunflowers * 25

        # 植物攻击
        for row in range(self.num_rows):
            for col in range(self.num_cols):
                # 豌豆射手攻击
                if self.lawn[row, col, 0] == 1:
                    for c in range(col + 1, self.num_cols):
                        zombie_idx = np.where(self.lawn[row, c, len(self.plants):] > 0)[0]
                        if len(zombie_idx) > 0:
                            self.lawn[row, c, len(self.plants) + zombie_idx[0]] = 0
                            reward += 140  # 增加三线射手的奖励
                            break
                
                # 寒冰射手攻击（减速效果）
                elif self.lawn[row, col, 3] == 1:
                    for c in range(col + 1, self.num_cols):
                        zombie_idx = np.where(self.lawn[row, c, len(self.plants):] > 0)[0]
                        if len(zombie_idx) > 0:
                            self.lawn[row, c, len(self.plants) + zombie_idx[0]] = 0
                            reward += 120
                            break
                
                # 双发射手攻击（可以攻击两个僵尸）
                elif self.lawn[row, col, 4] == 1:
                    zombies_hit = 0
                    for c in range(col + 1, self.num_cols):
                        zombie_idx = np.where(self.lawn[row, c, len(self.plants):] > 0)[0]
                        if len(zombie_idx) > 0:
                            self.lawn[row, c, len(self.plants) + zombie_idx[0]] = 0
                            reward += 100
                            zombies_hit += 1
                            if zombies_hit >= 2:
                                break
                
                # 大嘴花攻击（范围伤害）
                elif self.lawn[row, col, 5] == 1:
                    for r in [max(0, row-1), row, min(self.num_rows-1, row+1)]:
                        for c in range(col + 1, min(col + 3, self.num_cols)):
                            zombie_idx = np.where(self.lawn[r, c, len(self.plants):] > 0)[0]
                            if len(zombie_idx) > 0:
                                self.lawn[r, c, len(self.plants) + zombie_idx[0]] = 0
                                reward += 80
                
                # 食人花攻击（高伤害近战）
                elif self.lawn[row, col, 8] == 1:
                    if col < self.num_cols - 1:
                        zombie_idx = np.where(self.lawn[row, col + 1, len(self.plants):] > 0)[0]
                        if len(zombie_idx) > 0:
                            self.lawn[row, col + 1, len(self.plants) + zombie_idx[0]] = 0
                            reward += 150
                
                # 三线射手攻击（三行攻击）
                elif self.lawn[row, col, 9] == 1:
                    for r in [max(0, row-1), row, min(self.num_rows-1, row+1)]:
                        for c in range(col + 1, self.num_cols):
                            zombie_idx = np.where(self.lawn[r, c, len(self.plants):] > 0)[0]
                            if len(zombie_idx) > 0:
                                self.lawn[r, c, len(self.plants) + zombie_idx[0]] = 0
                                reward += 100
                                break

        # 僵尸移动和攻击
        for row in range(self.num_rows):
            for col in range(self.num_cols - 1, -1, -1):
                zombie_idx = np.where(self.lawn[row, col, len(self.plants):] > 0)[0]
                if len(zombie_idx) > 0:
                    if col == 0:
                        self.game_over = True
                        reward -= 1000
                    else:
                        # 检查前方是否有植物
                        if np.sum(self.lawn[row, col - 1, :len(self.plants)]) > 0:
                            # 攻击植物
                            plant_idx = np.where(self.lawn[row, col - 1, :len(self.plants)] > 0)[0][0]
                            # 如果是坚果墙，增加其防御能力
                            if plant_idx == 2:
                                if random.random() > 0.3:  # 70% 的概率不会被摧毁
                                    continue
                            self.lawn[row, col - 1, plant_idx] = 0
                        else:
                            # 向前移动
                            self.lawn[row, col - 1, len(self.plants) + zombie_idx[0]] = 1
                            self.lawn[row, col, len(self.plants) + zombie_idx[0]] = 0

        return reward

    def render(self, mode='human'):
        """渲染游戏状态"""
        if mode == 'human':
            print('\n' + '=' * 50)
            print(f'Sun: {self.sun}')
            print(f'Plants remaining: {self.plants_remaining}')
            print(f'Zombies remaining: {self.zombies_remaining}')
            print('\nLawn state:')
            for row in range(self.num_rows):
                for col in range(self.num_cols):
                    cell = self.lawn[row, col]
                    if np.sum(cell) == 0:
                        print('.', end=' ')
                    elif np.sum(cell[:len(self.plants)]) > 0:
                        plant_idx = np.where(cell[:len(self.plants)] > 0)[0][0]
                        print(self.plants[plant_idx][0][0], end=' ')
                    else:
                        zombie_idx = np.where(cell[len(self.plants):] > 0)[0][0]
                        print(self.zombies[zombie_idx][0][0], end=' ')
                print()
            print('=' * 50 + '\n')